import os.path
import re
import time
import cv2
from threading import Thread
from PySide2.QtWidgets import QMainWindow, QApplication, QFileDialog
from PySide2.QtGui import QImage, QIcon, QPixmap, QCloseEvent
from pygrabber.dshow_graph import FilterGraph
from src.ui.myWindow import Ui_MainWindow
import src.utils.openpose.pyopenpose as op
from src.utils.Calculation import calAngle


class MyWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        self.count = 0
        self.camera_id = 0
        self.mode_id = 0
        self.cap = None
        self.isOpen = False  # 摄像头开启与否的标识
        self.key_points = None  # 点位信息
        self.direction = True  # True为左,False为右
        self.thread = None
        self.flag_1 = False
        self.flag_2 = False
        self.flag_3 = False

        self.ui.button_open.clicked.connect(self.open)
        self.ui.button_exit.clicked.connect(self.exit)
        self.ui.button_clear.clicked.connect(self.clearLog)
        self.ui.button_save.clicked.connect(self.saveLog)
        self.ui.button_clear0.clicked.connect(self.dataReset)
        self.ui.button_changelr.clicked.connect(self.changeDirection)

        # 获取摄像头列表
        graph = FilterGraph()
        self.camera_list = graph.get_input_devices()
        self.ui.comboBox_camera.addItems(self.camera_list)
        self.ui.comboBox_camera.currentIndexChanged.connect(self.changeCamera)
        self.ui.comboBox_camera.setCurrentIndex(self.camera_id)

        # 设置模式选择框
        self.mode_data = {
            "name": ["肘部弯曲", "肩关节内收外展", "肩关节屈伸", "坐位伸膝", "直抬腿"],
            "benefit": ["增强肘关节运动能力，\n改善手部灵活性，锻炼\n肱二头肌、肱三头肌和\n前臂肌肉，增强骨骼和\n关节的健康", "", "", "", ""],
            "desc": ["将肘部伸直，屈肘90度\n伸展180度，重复屈伸\n\n20次/组", "", "", "", ""],
            "image": ["./assets/mode0.png", "./assets/mode1.png", "./assets/mode2.png", "./assets/mode3.png",
                      "./assets/mode4.png"],
            "count": [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]
        }
        self.ui.comboBox_mode.addItems(self.mode_data["name"])
        self.ui.comboBox_mode.currentIndexChanged.connect(self.changeMode)
        self.ui.comboBox_mode.setCurrentIndex(self.mode_id)
        self.showData()

    def open(self):
        """开启另一个线程,启动Openpose"""
        if self.isOpen:
            return
        self.thread = Thread(target=self.openPose)  # 开启另一个线程,启动openpose
        self.thread.start()
        self.isOpen = True  # 摄像头已开
        self.printLog('打开摄像头,启动OpenPose')

    def openPose(self):
        """打开Openpose"""
        params = {
            "model_folder": r"./utils/openpose/models/",
            "net_resolution": "160x80",
            "model_pose": "BODY_25",
            "frame_step": 2,
            "process_real_time": False,
            "face_detector": 0,
            "disable_blending": False
        }
        opWrapper = op.WrapperPython()
        opWrapper.configure(params)
        opWrapper.start()
        self.cap = cv2.VideoCapture(self.camera_id)
        while True:
            if not self.isOpen:
                self.key_points = None
                if self.flag_3:
                    self.ui.label_video.setStyleSheet('font-size:30px;')
                    self.ui.label_video.setText(f'恭喜您已完成{self.mode_data["name"][self.mode_id]}左右各一组，请稍作休息')
                    self.printLog(h1=f'{self.mode_data["name"][self.mode_id]}', text=' 已完成左右各 ',
                                  h2=f'{self.mode_data["count"][0][2]}组')
                    self.flag_3 = False
                else:
                    self.ui.label_video.setPixmap(QPixmap(""))
                break
            try:
                start_time = time.time()
                ret, frame = self.cap.read()
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                datum = op.Datum()
                datum.cvInputData = frame_rgb
                opWrapper.emplaceAndPop([datum])
                if datum.poseKeypoints.ndim != 0:  # 通过判断维度来检查是否检测到肢体
                    self.key_points = datum.poseKeypoints[0]
                    self.chooseMode()
                else:
                    self.key_points = None
                output_image = datum.cvOutputData
                end_time = time.time()
                seconds = end_time - start_time
                fps = round(1 / seconds)
                self.ui.label_fps.setText(f'{fps}')
                qImg = QImage(output_image.data, output_image.shape[1], output_image.shape[0], QImage.Format_RGB888)
                self.ui.label_video.setPixmap(QPixmap.fromImage(qImg))
            except Exception as e:
                print(e)
                break
            cv2.waitKey(0)

    def changeCamera(self):
        """切换摄像头"""
        if not self.isOpen:
            self.camera_id = self.ui.comboBox_camera.currentIndex()
        else:
            self.exit()
            self.camera_id = self.ui.comboBox_camera.currentIndex()
            self.open()
        self.printLog(text='切换摄像头为 ', h2=f'{self.camera_list[self.camera_id]}')

    def exit(self):
        """关闭Openpose,关闭摄像头"""
        if not self.isOpen:
            return
        self.isOpen = False
        self.cap.release()
        cv2.destroyAllWindows()
        self.printLog('关闭摄像头, 关闭OpenPose')

    def printLog(self, text, h1='', h2=''):
        """
        输出日志,并携带日志时间
        :param text: 普通日志内容
        :param h1:高亮日志内容1
        :param h2:高亮日志内容2
        """
        curTime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        self.ui.textBrowser.append(f'<p style="color:gray;">[{curTime}]</p>')
        self.ui.textBrowser.append(
            f'<span style="color: deepskyblue">{h1}</span>{text}<span style="color: deepskyblue">{h2}</span>')

    def showData(self):
        """展示模式名称以及该模式下的各种描述信息"""
        self.ui.label_modename.setText(self.mode_data["name"][self.mode_id])
        self.ui.label_benefit.setText(self.mode_data["benefit"][self.mode_id])
        self.ui.label_desc.setText(self.mode_data["desc"][self.mode_id])
        self.ui.label_sketchmap.setPixmap(QPixmap(self.mode_data["image"][self.mode_id]))
        self.ui.label_sketchmap.setMaximumSize(326, 174)
        self.ui.label_count.setText(f'次数:{self.mode_data["count"][self.mode_id][0 if self.direction else 1]}')

    def dataReset(self):
        """清零"""
        self.mode_data["count"][self.mode_id][0 if self.direction else 1] = 0
        self.ui.label_count.setText(f'次数:{0}')
        self.printLog(h1=f'{self.mode_data["name"][self.mode_id]}', text=' 的计数已清零')

    def clearLog(self):
        """清除日志"""
        self.ui.textBrowser.clear()

    def saveLog(self):
        """保存日志"""
        curTime = time.strftime("%Y-%m-%d %H-%M-%S", time.localtime())
        fileName, _ = QFileDialog.getSaveFileName(self, "保存日志", f"./log_{curTime}",
                                                  "Text Files (*.txt);;MarkDown Files (*.md);;HTML Files (*.html)")
        extName = os.path.splitext(fileName)[-1]
        if extName == ".txt":
            text = self.ui.textBrowser.toPlainText()
        else:
            listRe = re.findall(r'<p[\s\S]*?</p>', self.ui.textBrowser.toHtml())
            text = re.sub(r' style=" m.*?x;"', '', '\n'.join(listRe))
            if extName == ".html":
                fileBaseName = os.path.basename(fileName).split('.')[0]
                text = f'<!DOCTYPE html><html lang="zh"><head><meta charset="UTF-8"><title>{fileBaseName}</title></head><body>' + text + '</body></html>'
        with open(fileName, 'w', encoding='utf-8') as f:
            f.write(text)

    def chooseMode(self):
        """通过mode_id来判断该执行哪个检测函数"""
        if self.mode_id == 0:
            self.mode_0()
        elif self.mode_id == 1:
            self.mode_1()
        elif self.mode_id == 2:
            self.mode_2()
        elif self.mode_id == 3:
            self.mode_3()
        elif self.mode_id == 4:
            self.mode_4()

    def changeDirection(self):
        """改变左右"""
        self.direction = not self.direction
        self.flag_1 = False
        self.flag_2 = False
        self.ui.label_count.setText(f'次数:{self.mode_data["count"][self.mode_id][0 if self.direction else 1]}')
        self.ui.label_tip.setStyleSheet('font-size:25px;')
        self.ui.label_tip.setText('提示')
        self.ui.button_changelr.setText(f"切换({'左' if self.direction else '右'})")
        self.printLog("切换左右，当前为 ", h2=f"{'左' if self.direction else '右'}")

    def showComboBox(self):
        self.mode_id = self.ui.comboBox_mode.currentIndex()
        self.ui.button_changelr.setEnabled(True)
        self.ui.button_changelr.setText(f"切换({'左' if self.direction else '右'})")
        self.showData()

    def countdown(self):
        time_left = 3
        self.ui.label_video.setStyleSheet('font-size:60px;')
        while time_left > 0:
            self.ui.label_video.setText(f"自动切换至{'左' if self.direction else '右'}边,{time_left}...")
            time.sleep(1)
            time_left = time_left - 1

    def changeMode(self):
        """切换模式"""
        self.flag_1 = False
        self.flag_2 = False
        if not self.isOpen:
            self.showComboBox()
        else:
            self.exit()
            self.showComboBox()
            self.open()
        self.printLog(text='切换到 ', h2=f'{self.mode_data["name"][self.mode_id]}')

    def mode_0(self):
        """
        肘部弯曲训练
        动作要领：将肘部伸直，屈肘90度，伸展180度，重复屈伸。
        """
        dFlag = 0 if self.direction else 1
        point1 = self.key_points[5 - dFlag * 3]
        point2 = self.key_points[6 - dFlag * 3]
        point3 = self.key_points[7 - dFlag * 3]
        angle = calAngle(point1, point2, point3)
        self.ui.label_angle.setText(f"角度:{angle}")
        # if angle == 90:
        #     prtSc()
        if angle == "nan":
            return
        if self.flag_2:
            if self.flag_1:
                if 160 < angle < 180:
                    self.mode_data["count"][0][dFlag] += 1
                    self.flag_1 = False
                    self.flag_2 = False
                    self.ui.label_count.setText(f'次数:{self.mode_data["count"][0][dFlag]}')
                    if self.mode_data["count"][0][dFlag] == 20:
                        self.count += 1
                        self.dataReset()
                        self.changeDirection()
                        if self.count % 2 == 0:
                            self.mode_data["count"][0][2] += 1
                            self.flag_3 = True
                            self.count = 0
                            self.exit()
                            return
                        self.countdown()
                return
            if 80 < angle < 100:
                self.flag_1 = True
                self.ui.label_tip.setStyleSheet('font-size: 25px;color: yellow;')
                self.ui.label_tip.setText('肘部伸展')
            return
        if 160 < angle < 180:
            self.flag_2 = True
            self.ui.label_tip.setStyleSheet('font-size: 25px;color: deepskyblue;')
            self.ui.label_tip.setText('肘部屈曲')

    def mode_1(self):
        """肩关节内收外展"""
        dFlag = 0 if self.direction else 1
        point1 = self.key_points[1]
        point2 = self.key_points[5 - dFlag * 3]
        point3 = self.key_points[6 - dFlag * 3]
        angle = calAngle(point1, point2, point3)
        self.ui.label_angle.setText(f"角度:{angle}")
        if angle == "nan":
            return
        if self.flag_2:
            if self.flag_1:
                if 60 < angle < 80:
                    self.mode_data["count"][0][dFlag] += 1
                    self.flag_1 = False
                    self.flag_2 = False
                    self.ui.label_count.setText(f'次数:{self.mode_data["count"][0][dFlag]}')
                    if self.mode_data["count"][0][dFlag] == 20:
                        self.count += 1
                        self.dataReset()
                        self.changeDirection()
                        if self.count % 2 == 0:
                            self.mode_data["count"][0][2] += 1
                            self.flag_3 = True
                            self.count = 0
                            self.exit()
                            return
                        self.countdown()
                return
            if 160 < angle < 180:
                self.flag_1 = True
                self.ui.label_tip.setStyleSheet('font-size: 25px;color: yellow;')
                self.ui.label_tip.setText('肩关节内收')
            return
        if 60 < angle < 80:
            self.flag_2 = True
            self.ui.label_tip.setStyleSheet('font-size: 25px;color: deepskyblue;')
            self.ui.label_tip.setText('肩关节外展')

    def mode_2(self):
        """肩关节屈伸"""
        dFlag = 0 if self.direction else 1
        point1 = self.key_points[6 - dFlag * 3]
        point2 = self.key_points[1]
        point3 = self.key_points[8]
        angle = calAngle(point1, point2, point3)
        self.ui.label_angle.setText(f"角度:{angle}")
        if angle == "nan":
            return
        if self.flag_2:
            if self.flag_1:
                if 20 < angle < 40:
                    self.mode_data["count"][0][dFlag] += 1
                    self.flag_1 = False
                    self.flag_2 = False
                    self.ui.label_count.setText(f'次数:{self.mode_data["count"][0][dFlag]}')
                    if self.mode_data["count"][0][dFlag] == 20:
                        self.count += 1
                        self.dataReset()
                        self.changeDirection()
                        if self.count % 2 == 0:
                            self.mode_data["count"][0][2] += 1
                            self.flag_3 = True
                            self.count = 0
                            self.exit()
                            return
                        self.countdown()
                return
            if 70 < angle < 100:
                self.flag_1 = True
                self.ui.label_tip.setStyleSheet('font-size: 25px;color: yellow;')
                self.ui.label_tip.setText('肩关节后摆')
            return
        if 20 < angle < 40:
            self.flag_2 = True
            self.ui.label_tip.setStyleSheet('font-size: 25px;color: deepskyblue;')
            self.ui.label_tip.setText('肩关节上抬')

    def mode_3(self):
        """坐位伸膝"""
        dFlag = 0 if self.direction else 1
        point1 = self.key_points[12 - dFlag * 3]
        point2 = self.key_points[13 - dFlag * 3]
        point3 = self.key_points[14 - dFlag * 3]
        angle = calAngle(point1, point2, point3)
        self.ui.label_angle.setText(f"角度:{angle}")
        if angle == "nan":
            return
        if self.flag_2:
            if self.flag_1:
                if 70 < angle < 100:
                    self.mode_data["count"][0][dFlag] += 1
                    self.flag_1 = False
                    self.flag_2 = False
                    self.ui.label_count.setText(f'次数:{self.mode_data["count"][0][dFlag]}')
                    if self.mode_data["count"][0][dFlag] == 20:
                        self.count += 1
                        self.dataReset()
                        self.changeDirection()
                        if self.count % 2 == 0:
                            self.mode_data["count"][0][2] += 1
                            self.flag_3 = True
                            self.count = 0
                            self.exit()
                            return
                        self.countdown()
                return
            if 160 < angle < 180:
                self.flag_1 = True
                self.ui.label_tip.setStyleSheet('font-size: 25px;color: yellow;')
                self.ui.label_tip.setText('小腿放下')
            return
        if 70 < angle < 100:
            self.flag_2 = True
            self.ui.label_tip.setStyleSheet('font-size: 25px;color: deepskyblue;')
            self.ui.label_tip.setText('小腿上抬')

    def mode_4(self):
        dFlag = 0 if self.direction else 1
        point1 = self.key_points[1]
        point2 = self.key_points[12 - dFlag * 3]
        point3 = self.key_points[13 - dFlag * 3]
        angle = calAngle(point1, point2, point3)
        self.ui.label_angle.setText(f"角度:{angle}")
        if angle == "nan":
            return
        if self.flag_2:
            if self.flag_1:
                if 160 < angle < 180:
                    self.mode_data["count"][0][dFlag] += 1
                    self.flag_1 = False
                    self.flag_2 = False
                    self.ui.label_count.setText(f'次数:{self.mode_data["count"][0][dFlag]}')
                    if self.mode_data["count"][0][dFlag] == 20:
                        self.count += 1
                        self.dataReset()
                        self.changeDirection()
                        if self.count % 2 == 0:
                            self.mode_data["count"][0][2] += 1
                            self.flag_3 = True
                            self.count = 0
                            self.exit()
                            return
                        self.countdown()
                return
            if 100 < angle < 130:
                self.flag_1 = True
                self.ui.label_tip.setStyleSheet('font-size: 25px;color: yellow;')
                self.ui.label_tip.setText('腿放下')
            return
        if 160 < angle < 180:
            self.flag_2 = True
            self.ui.label_tip.setStyleSheet('font-size: 25px;color: deepskyblue;')
            self.ui.label_tip.setText('抬腿')

    def closeEvent(self, event: QCloseEvent):
        """在窗口关闭时关闭摄像头"""
        self.exit()


if __name__ == "__main__":
    app = QApplication([])
    app.setWindowIcon(QIcon("./assets/favicon.png"))
    myWin = MyWindow()
    myWin.show()
    app.exec_()
